import { fileURLToPath } from "url";
import { split, load } from "./loader";
import { config } from "dotenv";
import path from "path";
import fs from "fs";
import { genVectorStore, loadFromStore } from "./vectorStore";
import { useRetriever } from "./retriever";
import { RunnableSequence } from "@langchain/core/runnables";
import { convertDocsToString } from "./utils";
import { usePromptTemplate } from "./template";
import { useModel } from "../model";
import { StringOutputParser } from "@langchain/core/output_parsers";

/**
 * 根据本地的文档： 球状闪电
 */
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

async function main() {
  // 加载配置
  await config();
  // 加载本地文档
  const docs = await load(__dirname);
  // 文档分割
  const splitter = await split();
  const s_docs = await splitter.splitDocuments(docs);
  // 本地化存储
  const dir = path.resolve(__dirname, "..", "..", "data", "db", "qiu");
  // 判断是否存在这个文件夹目录
  if (!fs.existsSync(dir)) {
    console.log("向量存储目录不存在，正在创建...");
    await genVectorStore(s_docs, dir);
  }
  // 从本地加载 vectorStore
  const vectorStore = await loadFromStore(dir);

  const retriever = useRetriever(vectorStore);
  // const res = await retriever.invoke("原文中，谁提出了宏原子的假设？并详细介绍给我宏原子假设的理论")

  const contextRetriverChain = RunnableSequence.from([
    (input) => input.question,
    retriever,
    convertDocsToString,
  ]);

  // 真正的rag chain
  const ragChain = RunnableSequence.from([
    {
        context: contextRetriverChain,
        question: (input) => input.question,
    },
    usePromptTemplate(),
    useModel({}),
    new StringOutputParser(),
  ]);
  const answer = await ragChain.invoke({
    question: "用 200 字简述一下梗概"
  });  
  console.log(answer);
}

main();